import pandas as pd
import numpy as np
import matplotlib.pyplot as plt


def plot_item(path, item, annotation=None):
    data = pd.read_csv(path)
    # print(data['dLoss'])
    # pandas 数据可以直接观看其可视化形式
    data[item].plot(label=annotation)


item = 'RankLossBefore'
# plot_item(
#     './predatorprey/data/local/16agents/dicg2_ce_ppo_atype=general_res=True_entcoeff=0.1_grid=14_nagents=16_npreys=16_penalty=-1.50_stepcost=-0.10_avis=True_steps=200_nenvs=1_bs=6e+04_splits=3_miniepoch=10_seed=1/progress.csv',
#     item=item,
#     annotation='original_16agents'
# )
# plot_item(
#     './predatorprey/data/local/16agents/dicg2_ce_ppo_atype=general_res=True_entcoeff=0.1_grid=14_nagents=16_npreys=16_penalty=-1.50_stepcost=-0.10_avis=True_steps=200_nenvs=1_bs=6e+04_splits=3_miniepoch=10_seed=1_3_with_rank_loss/progress.csv',
#     item=item,
#     annotation='rank_loss_16agents'
# )
# plot_item(
#     './predatorprey/data/local/16agents/dicg2_ce_ppo_atype=general_res=True_entcoeff=0.1_grid=14_nagents=16_npreys=16_penalty=-1.50_stepcost=-0.10_avis=True_steps=200_nenvs=1_bs=6e+04_splits=3_miniepoch=10_seed=1_with_attention_score_rank_loss/progress.csv',
#     item=item,
#     annotation='attention_score_rank_loss_16agents'
# )
# plot_item(
#     './predatorprey/data/local/24agents/dicg2_ce_ppo_atype=general_res=True_entcoeff=0.1_grid=17_nagents=24_npreys=24_penalty=-1.50_stepcost=-0.10_avis=True_steps=200_nenvs=1_bs=6e+04_splits=3_miniepoch=10_seed=1/progress.csv',
#     item=item,
#     annotation='original_24agents'
# )
plot_item(
    './predatorprey/data/local/24agents/dicg2_ce_ppo_atype=general_res=True_entcoeff=0.1_grid=17_nagents=24_npreys=24_penalty=-1.50_stepcost=-0.10_avis=True_steps=200_nenvs=1_bs=6e+04_splits=3_miniepoch=10_seed=1_with_rank_loss/progress.csv',
    item=item,
    annotation='rank_loss_24agents'
)
plot_item(
    './predatorprey/data/local/24agents/dicg2_ce_ppo_atype=general_res=True_entcoeff=0.1_grid=17_nagents=24_npreys=24_penalty=-1.50_stepcost=-0.10_avis=True_steps=200_nenvs=1_bs=6e+04_splits=3_miniepoch=10_seed=1_with_rank_loss_2/progress.csv',
    item=item,
    annotation='rank_loss_24agents_repeat'
)
plot_item(
    './predatorprey/data/local/24agents/dicg2_ce_ppo_atype=general_res=True_entcoeff=0.1_grid=17_nagents=24_npreys=24_penalty=-1.50_stepcost=-0.10_avis=True_steps=200_nenvs=1_bs=6e+04_splits=3_miniepoch=10_seed=1_original_1500episodes/progress.csv',
    item=item,
    annotation='original_24agents_long'
)
# plot_item(
#     './predatorprey/data/local/32agents/dicg2_ce_ppo_atype=general_res=True_entcoeff=0.1_grid=20_nagents=32_npreys=32_penalty=-1.50_stepcost=-0.10_avis=True_steps=200_nenvs=1_bs=6e+04_splits=3_miniepoch=10_seed=1/progress.csv',
#     item=item,
#     annotation='original_32agents'
# )
# plot_item(
#     './predatorprey/data/local/32agents/dicg2_ce_ppo_atype=general_res=True_entcoeff=0.1_grid=20_nagents=32_npreys=32_penalty=-1.50_stepcost=-0.10_avis=True_steps=200_nenvs=1_bs=6e+04_splits=3_miniepoch=10_seed=1_1_with_rank_loss/progress.csv',
#     item=item,
#     annotation='rank_loss_32agents'
# )
plt.legend()
plt.title(item)
plt.show()
